#include <kutils.h>
#include "ata.h"
#include "irq.h"
#include "display.h"
#include "memory.h"

s32int wait_for_idle(ata_drive_t *d);
s32int wait_for_ready(ata_drive_t *d);
s32int wait_for_data(ata_drive_t *d);
s32int ata_wait_interrupt(ata_drive_t *d);
void ata_interrupt_handler(struct regs *r);
void ata_stall(void);
s32int ata_seek(ata_drive_t *d, u32int block);
u32int drive_exists(u32int driveID);
u32int check_partition(ata_drive_t *d, u32int pID);
void init_drive(u32int driveID);

ata_drive_t *drives[4];//4 drives each with 4 partitions
s8int d_exists[4];
volatile u32int ata_interrupted = 0;

s32int wait_for_idle(ata_drive_t *d)
{
	s32int time = ATA_TIMEOUT;
	//write_string("Waiting for idle...");
	for(time = ATA_TIMEOUT; time > 0; time--)
	{
		if(!(inportb(d->ports.status) & ATA_BUSY)) 
		{
			//write_string("[Drive is idle]\n");
			return 1;
		}
	}
	write_string("ATA Controller: Idle Timeout\n");
	return 0;
}

s32int wait_for_ready(ata_drive_t *d)
{
	s32int time = ATA_TIMEOUT;
	//write_string("Waiting for ready...");
	for(time = ATA_TIMEOUT; time > 0; time--)
	{
		if((inportb(d->ports.status) & ATA_READY)) 
		{
			//write_string("[Drive is ready]\n");
			return 1;
		}
	}
	write_string("ATA Controller: Ready Timeout\n");
	return 0;
}

s32int wait_for_data(ata_drive_t *d)
{
	s32int time = ATA_TIMEOUT;
	//write_string("Waiting for data...");
	for(time = ATA_TIMEOUT; time > 0; time--)
	{
		if((inportb(d->ports.status) & ATA_DATAREQ)) 
		{
			//write_string("[Drive has data]\n");
			return 1;
		}
	}
	write_string("ATA Controller: Data Timeout\n");
	return 0;
}

s32int ata_wait_interrupt(ata_drive_t *d)
{
	if(!d)
		return 0;

	//timeout is 1 second
	u32int timeout = 100;
	while(timeout)
	{
		if(!ata_interrupted)
		{
			timeout--;
		}
		else
		{
			ata_interrupted = 0;
			return 1;
		}
		wait(10);
	}
	ata_interrupted = 0;
	return 0;
}

void ata_interrupt_handler(struct regs *r)
{
	if(!r)
		return;
	//write_string("<hdint>");
	ata_interrupted = 1;
}

ata_drive_t *ata_get_drive(u32int driveID)
{
	if(driveID > 3)
		return 0;
	if(!d_exists[driveID])
	{
		return 0;
	}
	return drives[driveID];
}

void ata_stall(void)
{
	s32int i = 5;				// wait a s16int ammount of time
	while(i > 0) i--;
}

s32int ata_seek(ata_drive_t *d, u32int block)
{
	//null pos32inter
	if(d == 0)
		return 0;

	d->current_block = block;

	if(d->mode == ATA_LBA_MODE)
	{//lba addressing
		outportb(d->ports.LBA_low, (u8int)d->current_block);
  		outportb(d->ports.LBA_mid, (u8int)(d->current_block >> 8));
   	outportb(d->ports.LBA_high, (u8int)(d->current_block >> 16));
   	outportb(d->ports.drive_head, (u8int)(d->id | ((d->current_block >> 24) & 0x0F)));
	}
	else
	{//CHS addressing
		write_string(" chs ");
		u32int cyl = d->current_block / (d->heads * d->sectors);
		u32int tmp = d->current_block % (d->heads * d->sectors);
		u32int head = tmp / d->sectors;
		u32int sect = tmp % d->sectors + 1;

		outportb(d->ports.LBA_low, (u8int)sect);
  		outportb(d->ports.LBA_mid, (u8int)cyl);
   	outportb(d->ports.LBA_high, (u8int)(cyl >> 8));
   	outportb(d->ports.drive_head, (u8int)(d->id | head));
	}
	return 1;
}

s32int ata_read_partition(ata_partition_t *p, u8int *data, u32int sector, u32int offset, u32int length)
{
	if(!p->exists)
	{
		return 0;
	}
	return ata_read_drive(p->drive,data,p->lba_base + sector, offset, length);
}

s32int ata_write_drive(ata_drive_t *d, s8int *data, u32int sector, u32int offset, u32int length)
{
	if(!d)
		return 0;
	if(!data)
		return 0;
	if(sector > d->sectors)
		return 0;
	if(offset > 512)
		return 0;
	if(length > d->size)
		return 0;

	return 0;
}

s32int ata_read_drive(ata_drive_t *d, u8int *data, u32int sector, u32int offset, u32int length)
{
	if(!d){}
	if(!data){}
	if(sector > d->sectors){}
	if(offset > 512){}
	if(length > d->size){}
	
	if(!wait_for_idle(d))
	{
		return 0;
	}
	if(!ata_seek(d,sector))
	{
		return 0;
	}
	outportb(d->ports.sectors, 1);//read 1 sector
	wait(1);
	outportb(d->ports.command, ATA_READ);

	if(!ata_wait_interrupt(d))
	{
		write_string("ATA Controller: read stall\n");
	}

	if(!wait_for_idle(d))
	{
		return 0;
	}
	if(!wait_for_data(d))
	{
		return 0;
	}

  	u32int n = d->block_size / 2;
	u32int i = 0;
  	for(i = 0; i < n; i++)
	{
  		d->buffer[i] = (s16int)inportw(d->ports.data);
	}

	//for now limit the length to the block size
	if(length > d->block_size)
		length = d->block_size;

	memcpy(data,(s8int*)(d->buffer), (size_t)length);

	return 1;
}

u32int drive_exists(u32int driveID)
{
	switch(driveID)
	{
		case 0:
			//check controller zero
			outportb(0x1F3, 0x88);
			if(inportb(0x1F3) == 0x88)
			{//the controller is there, now check for drive zero
				outportb(0x1F6, 0 << 4 | 0xA0);//select drive 0
				wait(1);
				outportb(0x1F7, 0xEC); //tell the drive to identify itself
				
				s32int time = ATA_TIMEOUT;
				for(time = ATA_TIMEOUT; time > 0; time--)
				{
					if((inportb(0x1F7) & ATA_DATAREQ)) 
					{
						//the drive has data for us and therefore exists
						return 1;
					}
				}
				//timeout, the drive probably doesn't exist 
				//(or is so broken it shouldn't be used)
				return 0;
			}
		break;
		case 1:
			//check controller 0
			outportb(0x1F3, 0x88);
			if(inportb(0x1F3) == 0x88)
			{//the controller is there, now check for drive 1
				outportb(0x1F6, 1 << 4 | 0xB0);//select drive 1
				wait(1);
				outportb(0x1F7, 0xEC); //tell the drive to identify itself
				
				s32int time = ATA_TIMEOUT;
				for(time = ATA_TIMEOUT; time > 0; time--)
				{
					if((inportb(0x1F7) & ATA_DATAREQ)) 
					{
						//the drive has data for us and therefore exists
						return 1;
					}
				}
				//timeout, the drive probably doesn't exist 
				//(or is so broken it shouldn't be used)
				return 0;
			}
		break;
		case 2:
			//check controller 1
			outportb(0x173, 0x88);
			if(inportb(0x173) == 0x88)
			{//the controller is there, now check for drive 0
				outportb(0x176, 0 << 4 | 0xA0);//select drive 0
				wait(1);
				outportb(0x177, 0xEC); //tell the drive to identify itself
				
				s32int time = ATA_TIMEOUT;
				for(time = ATA_TIMEOUT; time > 0; time--)
				{
					if((inportb(0x177) & ATA_DATAREQ)) 
					{
						//the drive has data for us and therefore exists
						return 1;
					}
				}
				//timeout, the drive probably doesn't exist 
				//(or is so broken it shouldn't be used)
				return 0;
			}
		break;
		case 3:
			//check controller 1
			outportb(0x173, 0x88);
			if(inportb(0x173) == 0x88)
			{//the controller is there, now check for drive 
				outportb(0x176, 1 << 4 | 0xB0);//select drive 1
				wait(1);
				outportb(0x177, 0xEC); //tell the drive to identify itself
				
				s32int time = ATA_TIMEOUT;
				for(time = ATA_TIMEOUT; time > 0; time--)
				{
					if((inportb(0x177) & ATA_DATAREQ)) 
					{
						//the drive has data for us and therefore exists
						return 1;
					}
				}
				//timeout, the drive probably doesn't exist 
				//(or is so broken it shouldn't be used)
				return 0;
			}
		break;
		default:
			return 0;
		break;
	}
	return 0;
}

u32int check_partition(ata_drive_t *d, u32int pID)
{
	if(pID > 3)
	{
		return 0;
	}

	u32int entry_start = 0x1BE + (16 * pID);
	u8int *bs = (u8int*)kmalloc(512);
	ata_read_drive(d,bs,0,0,512);

	if(((bs[entry_start] & 0x80) != 0x80) & ((bs[entry_start] & 0x00) != 0x00))
	{//the partition is malformed. It doesnt exist.
		//write_string("\nMalformed Partition!\n");
		kfree(bs);
		return 0;
	}
	if(*(u32int*)(bs + entry_start + 0x0C) < 1)
	{//the partition is less than 1 sector large. It doesn't exist.
		//write_string("\nPartition size < 1 sector!\n");
		kfree(bs);
		return 0;
	}
	if(bs[entry_start + 0x04] < 1)
	{//bad partition type. It doesn't exist.
		//write_string("\nBad Partition Type!\n");
		kfree(bs);
		return 0;
	}

	//it exists. Now fill in the data
	d->partitions[pID].exists = 1;
	d->partitions[pID].drive = d;
	d->partitions[pID].status = bs[entry_start];
	d->partitions[pID].type = bs[entry_start + 0x04];
	d->partitions[pID].lba_base = *(u32int*)(bs + entry_start + 0x08);
	d->partitions[pID].lba_length = *(u32int*)(bs + entry_start + 0x0C);

	kfree(bs);
	return 1;
}

void init_drive(u32int driveID)
{
	u32int i = 0;

	ASSERT(driveID <= 3);	

	drives[driveID] = (ata_drive_t*)kmalloc(sizeof(ata_drive_t));
	ata_drive_t *d = drives[driveID];
	d->index = driveID;

	switch(driveID)
	{
		case 0:
			//known ports for the first ATA drive
			d->ports.data			 =	 0x1F0;
			d->ports.error			 =	 0x1F1;
			d->ports.features		 =	 0x1F1;
			d->ports.sectors		 =	 0x1F2;
			d->ports.LBA_low		 =	 0x1F3;
			d->ports.LBA_mid		 =	 0x1F4;
			d->ports.LBA_high		 =	 0x1F5;
			d->ports.drive_head	 =	 0x1F6;
			d->ports.status		 =	 0x1F7;
			d->ports.command		 =	 0x1F7;
			d->ports.alt_status	 =	 0x3F6;
			d->ports.control		 =	 0x3F6;
			d->id						 =  ATA_MASTER_CHS;
		break;
		case 1:
			//known ports for the second ATA drive
			d->ports.data			 =	 0x1F0;
			d->ports.error			 =	 0x1F1;
			d->ports.features		 =	 0x1F1;
			d->ports.sectors		 =	 0x1F2;
			d->ports.LBA_low		 =	 0x1F3;
			d->ports.LBA_mid		 =	 0x1F4;
			d->ports.LBA_high		 =	 0x1F5;
			d->ports.drive_head	 =	 0x1F6;
			d->ports.status		 =	 0x1F7;
			d->ports.command		 =	 0x1F7;
			d->ports.alt_status	 =	 0x3F6;
			d->ports.control		 =	 0x3F6;
			d->id						 =  ATA_SLAVE_CHS;
		break;
		case 2:
			//known ports for the third ATA drive
			d->ports.data			 =	 0x170;
			d->ports.error			 =	 0x171;
			d->ports.features		 =	 0x171;
			d->ports.sectors		 =	 0x172;
			d->ports.LBA_low		 =	 0x173;
			d->ports.LBA_mid		 =	 0x174;
			d->ports.LBA_high		 =	 0x175;
			d->ports.drive_head	 =	 0x176;
			d->ports.status		 =	 0x177;
			d->ports.command		 =	 0x177;
			d->ports.alt_status	 =	 0x376;
			d->ports.control		 =	 0x376;
			d->id						 =  ATA_MASTER_CHS;
		break;
		case 3:
			//known ports for the forth ATA drive
			d->ports.data			 =	 0x170;
			d->ports.error			 =	 0x171;
			d->ports.features		 =	 0x171;
			d->ports.sectors		 =	 0x172;
			d->ports.LBA_low		 =	 0x173;
			d->ports.LBA_mid		 =	 0x174;
			d->ports.LBA_high		 =	 0x175;
			d->ports.drive_head	 =	 0x176;
			d->ports.status		 =	 0x177;
			d->ports.command		 =	 0x177;
			d->ports.alt_status	 =	 0x376;
			d->ports.control		 =	 0x376;
			d->id						 =  ATA_SLAVE_CHS;
		break;
		default:

		break;
	}

	//write_string("Testing ATAs32interface...\n");

	//select the correct drive
	outportb(d->ports.drive_head, (u8int)d->id);
	wait(1);

	if(wait_for_idle(d) == -1)
	{
		d_exists[driveID] = 0;
		return;
	}
	if(wait_for_ready(d) == -1)
	{
		d_exists[driveID] = 0;
		return;
	}
		
	//select the drive again and tell it to identify itself
	outportb(d->ports.drive_head, (u8int)d->id);
	wait(10);
	outportb(d->ports.command, ATA_IDENTIFY);

	if(!ata_wait_interrupt(d))
	{
		write_string("ATA Controller: identify stall\n");
	}

	if(wait_for_idle(d) == -1)
	{
		d_exists[driveID] = 0;
		return;
	}
	if(wait_for_data(d) == -1)
	{
		d_exists[driveID] = 0;
		return;
	}

	for(i = 0; i<256; i++) 
	{
		d->buffer[i] = (s16int)inportw(d->ports.data);//read the data 
		//write_hex((u16int)d->buffer[i]); write_char(' ');
	}

	d->cylinders = *(u16int*)(d->buffer + 1);
  	d->heads = *(u16int*)(d->buffer + 3);
  	d->sectors = *(u16int*)(d->buffer + 6);

	//write_string(" C: ");
	//write_number(d->cylinders);
	//write_string(" H: ");
	//write_number(d->heads);
	//write_string(" S: ");
	//write_number(d->sectors);

	//supports LBA?
	if(d->buffer[49] & 0x200)
	{
		d->mode = ATA_LBA_MODE;
		//d->mode = ATA_CHS_MODE;
		if(d->id == ATA_MASTER_CHS)
		{
			//d->id = ATA_MASTER_CHS;
			d->id = ATA_MASTER_LBA_28;
		}
		if(d->id == ATA_SLAVE_CHS)
		{
			//d->id = ATA_SLAVE_CHS;
			d->id = ATA_SLAVE_LBA_28;
		}
	}
	else
	{
		write_string("ATA Controller: Drive ");
		write_number(driveID);
		write_string(" Doesn't support LBA!\n");
		d->mode = ATA_CHS_MODE;
		if(d->id == ATA_MASTER_CHS)
		{
			d->id = ATA_MASTER_CHS;
		}
		if(d->id == ATA_SLAVE_CHS)
		{
			d->id = ATA_SLAVE_CHS;
		}
	}

	d->blocks = d->cylinders * d->heads * d->sectors;

  	d->block_size = 512;//d->buffer[5];
  	d->size = (d->blocks * d->block_size);

	//convert the endianness of the strings
  	s8int tmp;		//temporary place to store a s8intacter
  	s8int *strptr;	//will pos32int to the currently working memory position
  
  	strptr = (s8int*)d->buffer;//start by pos32inting to the beginning of the buffer
  	for(i=0; i<512; i+=2)
  	{
  		tmp = *strptr;				//store the current s8intacter
  		*strptr = *(strptr+1);	//move the next s8intacter s32into the current position
  		strptr++;					//go to the next position
  		*strptr = tmp;				//put the first s8intacter there
  		strptr++;					//go to the next position
  	}

	memcpy(d->serial_number,(s8int *)(d->buffer+10), 20);
  	memcpy(d->firmware_revision, (s8int *)(d->buffer+23), 8);
  	memcpy(d->model, (s8int *)(d->buffer+27),40);

	i = 0;
	for(i = 0; i < 4; i++)
	{
		check_partition(d,i);
	}
}

void init_ata(void)
{
#ifdef FLUIDIUM_DEBUG
	write_string("ATA Controller: Init\n");
#endif
	irq_install_handler(14, ata_interrupt_handler);

	u32int i = 0;	
	for(i = 0; i < 4; i++)
	{
		if(!drive_exists(i))
		{
			//write_string("\nDrive ");
			//write_number(i);
			//write_string(" was not detected.\n");
			d_exists[i] = 0;
		}
		else
		{
			//write_string("\nDrive ");
			//write_number(i);
			//write_string(" exists.\n");
			//say the drive exists, this could be overwritten if something goes wrong.
			d_exists[i] = 1;
			init_drive(i);
		}
	}
}

/*
void ata_dump_drive_info(window *win, ata_drive_t *d)
{
	write_window_string(win,"\n");
	if(d == 0)
	{
		write_window_string(win,"Drive "); 
		write_window_number(win,d->index); 
		write_window_string(win," does not exist.\n");
		write_window_string(win,"\n");
		return;
	}

	write_window_string(win,"  Model: ");
	write_window_string(win,d->model);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Serial Number: ");
	write_window_string(win,d->serial_number);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Firmware Revision: ");
	write_window_string(win,d->firmware_revision);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Size (kilobytes): ");
	write_window_number(win,d->size / 1024);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Size (megabytes): ");
	write_window_number(win,d->size / 1024 / 1024);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Mode: ");
	write_window_number(win,d->mode);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Sectors: ");
	write_window_number(win,d->sectors);
	write_window_string(win," \n");

	//the first partition
	write_window_s8int(win,'\n');
	write_window_string(win,"  First Partition:");
	write_window_s8int(win,'\n');
	write_window_string(win,"  Exists: ");
	write_window_number(win,d->partitions[0].exists);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Status: ");
	write_window_hex(win,d->partitions[0].status);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Type: ");
	write_window_hex(win,d->partitions[0].type);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Base Sector: ");
	write_window_number(win,d->partitions[0].lba_base);
	write_window_s8int(win,'\n');
	write_window_string(win,"  Length (Sectors): ");
	write_window_number(win,d->partitions[0].lba_length);
	write_window_s8int(win,'\n');

	write_window_string(win,"\n");
}*/
